import cv2
import numpy as np
import os
import pydicom as dicom
import pandas as pd
import shutil
import SimpleITK as sitk
from pathlib import Path
from tqdm import tqdm

annotation_Box_path = "/data1/home/liukai/AllData/Duke/Duke/Annotation_Boxes.xlsx"


# 获取dicom文件的pixel_array
def get_pixel_array(dicom_file_path):
    dcm_ori = dicom.dcmread(dicom_file_path)
    arr = dcm_ori.pixel_array
    arr = arr.astype(np.int32)
    return arr


# 最大最小归一化
def pixel_array_normalize(pixel_array):
    return 255.0 * (pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min())


def get_subtraction(array1, array2):
    """
    获取剪影图像
    :param dicom1:  增强后的dicom文件路径
    :param dicom2:  增强前的dicom文件路径
    :return:        剪影图像矩阵
    """
    pa1 = array1
    pa2 = array2
    pa1 = pixel_array_normalize(pa1)
    pa2 = pixel_array_normalize(pa2)
    img = pa1 - pa2
    del pa1
    del pa2
    # 截断负值
    img = np.clip(img, a_min = 0, a_max = img.max())
    img = pixel_array_normalize(img)
    # 计算直方图
    img = np.uint8(img)
    hist = cv2.calcHist([img], [0], None, [255], [1, 256])
    # plt.plot(hist)
    # plt.show()
    hist_sum = np.cumsum(hist)
    # 截断归一化
    min_val = 0
    # min_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.01, side='left')
    max_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.9998, side = 'left')
    img = np.clip(img, a_min = min_val, a_max = max_val)
    img = 255.0 * (img - min_val) / (max_val - min_val)

    return img


def dukeFileChoose(dir_path):
    '''
    即筛选出的脂肪饱和处理前序列文件名称列表、第二次增强后的文件名称列表
    :param dir_path: duke原数据的文件夹路径
    :return: dict{'pre':[],'post':[]}
    '''

    def judge_filename(filename):
        '''
        最终筛选出917个例子，3D形状从92 ~256

        用于下面特定的if条件，筛选出有ax 3d dyn且不包含['2nd','3rd','4th','1st']的文件名
        :param filename:
        :return:

        '''
        a = ['2nd', '3rd', '4th', '1st']
        for s in a:
            if s in filename:
                return True
        return False

    def judge_filename2(filename):
        '''
        用于下面特定的if条件，筛选出有ax 3d dyn且不包含['2nd','3rd','4th','1st']的文件名
        :param filename:
        :return:
        '''
        a = ['']
        for s in a:
            if s in filename:
                return True
        return False

    global annotation_Box_path

    count = 0
    no_pre_count = 0
    pre_path_list = []
    post_path_list = []
    case_len_list = []  # 记录筛选出来的每个case的dcm文件数量
    for case in os.listdir(dir_path):
        if case == "Breast_MRI_151":  # 特殊处理，这个的pre序列和第二次增强对比序列的维度不一样
            continue
        b = os.listdir(os.path.join(dir_path, case))
        for tmp_dir in os.listdir(os.path.join(dir_path, case)):
            pre_sequence_dir = None
            second_enhanse_dir = None
            only_one = 0
            file_list = os.listdir(os.path.join(dir_path, case, tmp_dir))
            okList = []
            # -----找出对比前的影像-------#
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                # if  'pre' in sequence_dir \
                #         or '-ax 3d dyn' in sequence_dir or '-ax dynamic' in sequence_dir \
                #         or '-ax dyn' in sequence_dir:
                if '-Ax Vibrant' in sequence_dir \
                        or ('pre' in sequence_dir and 't1 pre' not in sequence_dir) \
                        or ('-ax 3d dyn' in sequence_dir and not judge_filename(sequence_dir)) \
                        or '-ax dynamic' in sequence_dir \
                        or '-ax dyn-' in sequence_dir or \
                        '-ax 3d dyn pre' in sequence_dir:
                    # if 'pre' in sequence_dir:
                    # if '-ax 3d dyn' in sequence_dir:
                    pre_sequence_dir = sequence_dir
                    only_one += 1
                    okList.append(sequence_dir)
            if only_one != 1:
                if only_one == 0:
                    no_pre_count += 1
                    print("没找到对比前序列", os.listdir(os.path.join(dir_path, case, tmp_dir)))
                else:
                    # print("有多个文件符合条件，无法确定哪个是对比前序列")
                    # print(case,okList)
                    pass
                continue
            # ------------找出第二次增强对比后的影像-------------#
            only_one = 0
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                if '2nd' in sequence_dir or 'Ph2Ax' in sequence_dir or '2ax' in sequence_dir:
                    only_one += 1
                    second_enhanse_dir = sequence_dir
            if only_one == 1:  # 找到符合条件的对比前序列和第二次增强的对比后序列

                count += 1
            else:
                if only_one == 0:
                    print("没找到第二次增强序列")
                else:
                    print("有多个文件符合条件，无法确定哪个是第二次增强序列")
            # --处理筛选出来的数据--#
            post_path = os.path.join(dir_path, case, tmp_dir, second_enhanse_dir)
            pre_path = os.path.join(dir_path, case, tmp_dir, pre_sequence_dir)
            pre_path_list.append(pre_path)
            post_path_list.append(post_path)

            # 对比组数据的dcm文件数量是否一致（也就是对比维度是否一致）
            # reader = sitk.ImageSeriesReader()
            # dicom_names = reader.GetGDCMSeriesFileNames(pre_path)
            # reader.SetFileNames(dicom_names)
            # image = reader.Execute()
            # image_array=sitk.GetArrayFromImage(image)
            # print(pre_path,image_array.shape)

    # print('序列数：', count)
    # print("筛选出来的数据的维度的范围",case_len_list[0],case_len_list[-1])
    return {"pre": pre_path_list, "post": post_path_list}








def check_nii_dcm(file_list):
    '''
    :param file_list: dcm文件夹列表
    :param nii_save_dir: nii文件保存的地方
    :return: None
    '''

    down_count = 0
    up_count = 0
    x_list = []
    y_list = []
    z_list = []
    count = 0
    for file in tqdm(file_list):
        dcm_file_list = os.listdir(file)
        dcm_file_list = sorted(dcm_file_list)  # 一定要为文件名排个序
        dcm_first = os.path.join(file, dcm_file_list[0])
        dcm_ori_first = dicom.dcmread(dcm_first)
        z_first = dcm_ori_first.ImagePositionPatient[-1]
        dcm_last = os.path.join(file, dcm_file_list[-1])
        dcm_ori_last = dicom.dcmread(dcm_last)
        z_last = dcm_ori_last.ImagePositionPatient[-1]
        is_revers = False
        if z_first < z_last:  # 正常的升序
            up_count += 1
        elif z_first > z_last:
            is_revers = True
            down_count += 1
        is_revers = (not is_revers)  # 跟上次实验反过来试试
        dicom_files = dcm_file_list  # 获取目录下所有 .dcm 文件

        slices = [dicom.read_file(os.path.join(file, dicom_file), force = True) for dicom_file in dicom_files]
        if is_revers:
            slices.reverse()
        # ----看x和y轴的spacing是否相同---#
        slice_x = [slice.PixelSpacing[0] for slice in slices]
        slice_y = [slice.PixelSpacing[0] for slice in slices]
        slice_x_set = set(slice_x)
        slice_y_set = set(slice_y)
        assert len(slice_x_set) == 1 and len(slice_y_set) == 1
        spacing_x = slice_x[0]
        spacing_y = slice_y[0]
        # 计算z轴的spacing
        # 获取整个 CT 序列的 spacing, 等同于 SimpleITK 的 GetSpacing() 方法
        spacing_z_list = []
        for i in range(1, len(slices)):
            spacing_z_list.append(slices[i].SliceLocation - slices[i - 1].SliceLocation)
        spacing_z = abs(sum(spacing_z_list) / len(spacing_z_list))  # 获取Z轴的spacing
        count += 1
        # 获取全部dcm的pixel，组成一个3D数组，并保存为nii文件
        try:
            if spacing_z != 0:
                nii_save_dir='/data1/home/liukai/AllData/Duke/Duke/nii/pre'
                slices_pixel_list = [slice.pixel_array.astype(np.int32) for slice in slices]
                nii_file_name = os.path.dirname(file).split('/')[-2] + '.nii.gz'
                nii_file = sitk.ReadImage(os.path.join(nii_save_dir, nii_file_name))
                nii_file_array=sitk.GetArrayFromImage(nii_file)
                for i in range(len(nii_file_array)):
                    a=(nii_file_array[i]==slices_pixel_list[i]).all()
                    if a!=True:
                        print(nii_file_name,is_revers,"different!")
                        break


        except:
            print("warning z_spacing=0")
        finally:
            pass


# 获取duke数据集的剪影nii图像
def duke2subtraction(file):
    substract_path = "/data1/home/liukai/AllData/Duke/Duke/nii/substractSecond"
    nii_path = "/data1/home/liukai/AllData/Duke/Duke/nii"

    pre_dir = os.path.join(nii_path, "pre")
    post2_dir = os.path.join(nii_path, "post2")

    pre_nii = sitk.ReadImage(os.path.join(pre_dir, file))
    pre_array = sitk.GetArrayFromImage(pre_nii).astype(np.int32)

    post2_nii = sitk.ReadImage(os.path.join(post2_dir, file))
    post2_array = sitk.GetArrayFromImage(post2_nii).astype(np.int32)

    img_array = get_subtraction(pre_array, post2_array)
    spacing = pre_nii.GetSpacing()
    nii_file = sitk.GetImageFromArray(img_array)
    nii_file.SetSpacing(spacing)
    nii_file_name = file

    sitk.WriteImage(nii_file, os.path.join(substract_path, nii_file_name))
    print(file)


def get_center_slice(annotation_path):
    """
    获取每个病例病灶中心切片号
    :param annotation_path: 标注文件路径
    :return: 字典 key为病例ID value为切片号
    """
    annotations = pd.read_excel(annotation_path)
    caseId2sliceId = {}
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_slice = annotation[6]
        end_slice = annotation[7]
        caseId2sliceId[case_id] = int((end_slice - start_slice) / 2 + start_slice)
    return caseId2sliceId

# 获取duke数据集的剪影图像
def checkSubtraction(dir_path, save_path):
    caseId2sliceId = get_center_slice('/data1/home/liukai/AllData/Duke/Duke/Annotation_Boxes.xlsx')
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    L=dukeFileChoose(dir_path =dir_path)
    pre_file_list,post_file_list=L['pre'],L['post']
    count=0
    for i in range(len(pre_file_list)):
        #--------由dcm切片计算得到-------#
        case=pre_file_list[i].split('/')[-3]

        selectSliceId = caseId2sliceId[case]
        sliceIds=[selectSliceId]
        for sliceId in sliceIds:
            if sliceId < 100:
                sliceId = '0' + str(sliceId)
            else:
                sliceId = str(sliceId)

            dicom_file = '1-' + sliceId + '.dcm'
            post_path = os.path.join(post_file_list[i], dicom_file)
            pre_path = os.path.join(pre_file_list[i], dicom_file)
            pre_dcm_array=get_pixel_array(pre_path)
            post_dcm_array = get_pixel_array(post_path)
            img = get_subtraction(post_path, pre_path)
            count+=1
            print(os.path.join(save_path, case + '_second_subtraction_norm_' + dicom_file[2:5] + '.jpg'))






if __name__ == '__main__':
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # get_center_slice(annotation_path)

    # 画框
    # save_path = 'E:\\MRI\\Duke\\tmp'
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # draw_box(save_path, annotation_path)

    # 获取duke数据集的subtraction图像
    # dir_path = 'E:\\MRI\\Duke\\manifest-1680071275430\\Duke-Breast-Cancer-MRI'
    # save_path = 'E:\\MRI\\Duke\\subtractions'
    # duke2subtraction(dir_path, save_path)

    # 切片根据是否含病灶进行分类
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # save_path = 'E:\\MRI\\Duke'
    # dir_path = 'E:\\MRI\\Duke\\case'
    # duke_classify(annotation_path, dir_path, save_path)

    # --------单线程将dcm转化为nii--------#
    # dir_path = '/data1/home/liukai/AllData/Duke/Duke/Duke-Breast-Cancer-MRI'
    # pre_save_path = '/data1/home/liukai/AllData/Duke/Duke/nii/pre'
    # post2_save_path = '/data1/home/liukai/AllData/Duke/Duke/nii/post2'
    #
    # if not os.path.exists(pre_save_path):
    #     os.makedirs(pre_save_path)
    # if not os.path.exists(post2_save_path):
    #     os.makedirs(post2_save_path)
    # import time
    #
    # start = time.time()
    #
    # all_dcm_to_nii(dukeFileChoose(dir_path)['pre'],nii_save_dir = pre_save_path)
    # all_dcm_to_nii(dukeFileChoose(dir_path)['post'], nii_save_dir = post2_save_path)
    # end = time.time()
    #
    # print("ok!cost {}/min".format((end - start) / 60))

    # -----多线程将dcm转化为nii，这里需要执行两次（两次分别要修改下面的dukeFileChoose(dir_path)['pre']和函数里的路径参数）----#
    # import time
    # start=time.time()
    # pool_for_test = Pool(1)
    # result = pool_for_test.map(dcm_to_nii, dukeFileChoose(dir_path)['pre'])
    # end=time.time()
    #
    # print("ok!cost {}/min".format((end-start)/60))

    # ------------给dcm文件统一命名为三位数字--------#
    dir_path = '/data1/home/liukai/AllData/Duke/Duke/Duke-Breast-Cancer-MRI'
    check_nii_dcm(dukeFileChoose(dir_path)['pre'])